import asyncio
import collections
import itertools
import multiprocessing
import os
import random
import struct
import sys
import time
import queue
from typing import Optional
import logging
logging.basicConfig(level=logging.DEBUG)

import networkx
import matplotlib.pyplot as plt

transport_logger = logging.getLogger("transport")
transport_logger.setLevel(logging.INFO)
message_logger = logging.getLogger("message")
algorithm_logger = logging.getLogger("algorithm")
processing_logger = logging.getLogger("processing")


YIELD_TIMEOUT = 1000000 # ns == 1 ms
NEIGHBOUR_TIMEOUT = 3000000000 # ns = 2000ms
SENDER_OPTIONS = [0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF]
BUTTON_DEBOUNCE = 3000000000 # ns = 3000ms

BUTTON_MAP = { # (node, logical): phy
    (1, 0): 4, (1, 1): 0, (1, 2): 3, (1, 3): 2, (1, 4): 1, (1, 5): 5,
    (2, 0): 2, (2, 1): 3, (2, 2): 0, (2, 3): 1, (2, 4): 5, (2, 5): 4,
    (3, 0): 4, (3, 1): 0, (3, 2): 2, (3, 3): 1, (3, 4): 3, (3, 5): 5,
    (4, 0): 0, (4, 1): 5, (4, 2): 2, (4, 3): 4, (4, 4): 1, (4, 5): 3,
    (5, 0): 2, (5, 1): 3, (5, 2): 1, (5, 3): 0, (5, 4): 5, (5, 5): 4,
    (6, 0): 3, (6, 1): 5, (6, 2): 2, (6, 3): 1, (6, 4): 0, (6, 5): 4,
    (7, 0): 1, (7, 1): 4, (7, 2): 0, (7, 3): 3, (7, 4): 2, (7, 5): 5,
    (8, 0): 1, (8, 1): 0, (8, 2): 2, (8, 3): 3, (8, 4): 4, (8, 5): 5,
}
LED_MAP = { # (node, logical): phy
    (1, 0): 4, (1, 1): 3, (1, 2): 5, (1, 3): 1, (1, 4): 2, (1, 5): 0,
    (2, 0): 5, (2, 1): 3, (2, 2): 4, (2, 3): 2, (2, 4): 0, (2, 5): 1,
    (3, 0): 4, (3, 1): 3, (3, 2): 5, (3, 3): 1, (3, 4): 2, (3, 5): 0,
    (4, 0): 4, (4, 1): 2, (4, 2): 3, (4, 3): 0, (4, 4): 5, (4, 5): 1,
    (5, 0): 2, (5, 1): 4, (5, 2): 0, (5, 3): 5, (5, 4): 1, (5, 5): 3,
    (6, 0): 4, (6, 1): 0, (6, 2): 2, (6, 3): 5, (6, 4): 3, (6, 5): 1,
    (7, 0): 4, (7, 1): 0, (7, 2): 5, (7, 3): 2, (7, 4): 3, (7, 5): 1,
    (8, 0): 4, (8, 1): 5, (8, 2): 2, (8, 3): 0, (8, 4): 3, (8, 5): 1,
}

last_sender_by_led = {}
last_led_by_sender = {}

def prepare_packet(d):
    return b"\xFF\xFF\xFF\xFF\x00\x00\x00\x00" + struct.pack("<I", d)


def _get_next_sender(message, sender_index):
    if message & 0b1000 and not message & 0b100:
        recipient = (message & 0x00F00000) >> 20
        index = (message & 0x000F0000) >> 16
        last_sender = last_sender_by_led.get((recipient, index), None)
        if last_sender is not None and last_led_by_sender.get(last_sender, (None, None, False))[:2] == (recipient, index):
            transport_logger.debug("reusing led sender %d for %d,%d", last_sender, recipient, index)
            last_led_by_sender[last_sender] = (*last_led_by_sender[last_sender][:2], False)
            return last_sender, sender_index

    while not last_led_by_sender.get(SENDER_OPTIONS[sender_index], (0, 0, True))[2]:
        last_led_by_sender[SENDER_OPTIONS[sender_index]] = (*last_led_by_sender[SENDER_OPTIONS[sender_index]][:2], True)
        sender_index = (sender_index + 1) % len(SENDER_OPTIONS)

    if message & 0b1000 and not message & 0b100:
        transport_logger.debug("using new led sender %d for %d,%d", SENDER_OPTIONS[sender_index], recipient, index)
        last_led_by_sender[SENDER_OPTIONS[sender_index]] = (recipient, index, False)
        last_sender_by_led[recipient, index] = SENDER_OPTIONS[sender_index]
    else:
        last_led_by_sender.pop(SENDER_OPTIONS[sender_index], None)
    return SENDER_OPTIONS[sender_index], (sender_index + 1) % len(SENDER_OPTIONS)

async def transceive_loop(raw_in_queue, raw_out_queue, data_in_queue, data_out_queue):
    last_inbound = time.monotonic_ns()
    sender_index = 0
    while True:
        p = b""
        while len(p) < 4:
            c = await raw_in_queue.get()
            p += c
        n, = struct.unpack("<I", p)
        transport_logger.debug("<%#010x", n)
        log_message(n)
        parity = n.bit_count() % 2
        if parity:
            # successfully yielded control or just reinitialised
            last_inbound = time.monotonic_ns()
        if parity:
            await data_in_queue.put(n)
        if n == 0 or (n & 2 and parity):
            forwardable = True
            try:
                msg = data_out_queue.get_nowait()
            except asyncio.QueueEmpty:
                msg = 0x00000002
                forwardable = False
            if last_inbound < time.monotonic_ns() - YIELD_TIMEOUT // 2:
                msg |= 0x00000002
            msg &= ~0x0F000000
            if forwardable:
                sender, sender_index = _get_next_sender(msg, sender_index)
                msg |= sender << 24
            if msg.bit_count() % 2 == 0:
                msg ^= 0b1
            transport_logger.debug(">%#010x", msg)
            log_message(msg)
            raw_out_queue.put_nowait(prepare_packet(msg))

def log_message(n):
    message_type = n & 0b1100
    link_sender = n >> 28
    remote_sender = (n >> 24) & 0xF
    if not remote_sender:
        return
    if message_type == 0b0000:
        # update neighbours
        neighbours = ((n >> 4) & 0xF, (n >> 8) & 0xF, (n >> 12) & 0xF, (n >> 16) & 0xF, (n >> 20) & 0xF)
        message_logger.debug("%x: %x: neighbours: %x, %x, %x, %x, %x", link_sender, remote_sender, *neighbours)
    elif message_type == 0b0100:
        buttons = ((n >> 16) & 0b1, (n >> 17) & 0b1, (n >> 18) & 0b1, (n >> 19) & 0b1, (n >> 20) & 0b1, (n >> 21) & 0b1,
                   (n >> 22) & 0b1, (n >> 23) & 0b1)
        message_logger.debug("%x: %x: buttons: %x, %x, %x, %x, %x, %x (%x, %x)", link_sender, remote_sender, *buttons)
    else:
        recipient = (n >> 20) & 0xF
        led = (n >> 16) & 0xF
        rgb = (n >> 4) & 0xFFF
        reset = n & 0b0100
        message_logger.debug("%x: %x -> %x: led %d := %03X%s", link_sender, remote_sender, recipient, led, rgb, " and blank others" if reset else "")




neighbours_state = networkx.DiGraph()
connect_with_buttons = False
# edges contain a key of each node, which contains the logical pin number, as well as a key "colour"
buttons_graph = networkx.Graph()
EDGE_COLOURS = frozenset({0xF0F, 0xF00, 0x0F0, 0x00F, 0xFF0, 0x0FF, 0xFFF, 0x8FF})
pairing_node_pin: tuple[int, int, int, asyncio.Task] | None = None

async def update_neighbours(link_sender, remote_sender, neighbours):
    current_time = time.monotonic_ns()
    try:
        neighbours_state.remove_edges_from(list(neighbours_state.out_edges(remote_sender)))
    except networkx.NetworkXError:
        pass
    for i, neighbour in enumerate(neighbours):
        if neighbour != 0 and neighbour != 0xF:
            neighbours_state.add_edge(remote_sender, neighbour, time=current_time, pin=i)
    neighbours_state.add_edge(0, link_sender, time=current_time)

neighbour_leds = {}
async def update_neighbour_leds_loop(data_out_queue):
    while True:
        if connect_with_buttons:
            await asyncio.sleep(1)
            continue
        graph = get_graph()
        possible_colours = set(EDGE_COLOURS)
        for c, _ in neighbour_leds.values():
            possible_colours.remove(c)
        target = set()
        current = set(neighbour_leds.keys())
        for u, v, d in graph.edges(data=True):
            k = (u, d[u]["pin"], v, d[v]["pin"])
            target.add(k)
        for k in current - target:
            del neighbour_leds[k]
            data_out_queue.put_nowait(set_led(k[0], k[1] + 1, 0))
            data_out_queue.put_nowait(set_led(k[2], k[3] + 1, 0))
        for k in target - current:
            c = possible_colours.pop()
            neighbour_leds[k] = c, time.monotonic_ns()
            data_out_queue.put_nowait(set_led(k[0], k[1] + 1, c))
            data_out_queue.put_nowait(set_led(k[2], k[3] + 1, c))
        await asyncio.sleep(0.1)


async def refresh_neighbour_leds_loop(data_out_queue):
    while True:
        await asyncio.sleep(2)
        if connect_with_buttons:
            continue
        if data_out_queue.qsize() > 3:
            continue
        if not neighbour_leds:
            continue
        # refresh something just for the sake of it
        # could use a heap but this is easier
        oldest = min(neighbour_leds.items(), key=lambda x: x[1][1])
        processing_logger.debug("refreshing %d:%d from %d", oldest[0][0], oldest[0][1], oldest[1][1])
        data_out_queue.put_nowait(set_led(oldest[0][2], oldest[0][3] + 1, oldest[1][0]))
        neighbour_leds[oldest[0]] = oldest[1][0], time.monotonic_ns()


async def send_neighbours_mode(data_out_queue):
    while True:
        data_out_queue.put_nowait(0xFF060008 | (not connect_with_buttons) << 4)
        await asyncio.sleep(10)


def get_active_neighbours(node):
    try:
        return get_graph().adj[node]
    except KeyError:
        return []

def set_led(recipient, led, colour):
    return 0xFF000008 | recipient << 20 | colour << 4 | LED_MAP[recipient, led] << 16




async def stub_algorithm(start_node, data_out_queue):
    neighbours = get_active_neighbours(start_node)
    data_out_queue.put_nowait(set_led(start_node, 0, 0x800))
    for neighbour in neighbours:
        data_out_queue.put_nowait(set_led(neighbour, 0, 0x0F0))
    await asyncio.sleep(3)
    for neighbour in neighbours:
        data_out_queue.put_nowait(set_led(neighbour, 0, 0))
    data_out_queue.put_nowait(set_led(start_node, 0, 0))


def get_graph():
    if connect_with_buttons:
        return buttons_graph
    current_time = time.monotonic_ns()
    graph = networkx.Graph()
    graph.add_nodes_from(neighbours_state)
    try:
        graph.remove_node(0)
    except networkx.NetworkXError:
        assert not graph.nodes
        return graph
    for u, v, d in neighbours_state.out_edges.data():
        if u != 0:
            graph.add_edge(u, v)
            graph.edges[u, v][u] = d
    for u, v, d in list(graph.edges.data()):
        # ignore if one of the nodes reports the link doesn't exist, or if neither node has checked in recently
        if len(d.keys()) != 2 or (d[u]["time"] < current_time - NEIGHBOUR_TIMEOUT and d[v]["time"] < current_time - NEIGHBOUR_TIMEOUT):
            graph.remove_edge(u, v)
    return graph.subgraph(itertools.chain.from_iterable(networkx.node_connected_component(graph, n) for n, d in neighbours_state.adj[0].items() if d["time"] > current_time - NEIGHBOUR_TIMEOUT))

target_node = 8

def algorithm_control_filter(msg):
    if msg & 0b1100 == 0b0100:
        return True
    return False


async def bfs_algorithm(start_node, data_out_queue):
    graph = get_graph()
    pending = collections.deque(({start_node},))
    visited = {start_node}
    guess = None

    while pending and pending[0]:
        if guess in pending[0]:
            node = guess
        else:
            node = random.choice(tuple(pending[0]))
        pending[0].remove(node)
        if not pending[0]:
            pending.popleft()
        algorithm_logger.info("visit %d", node)
        if node == target_node:
            data_out_queue.put_nowait(set_led(node, 0, 0x0F0))
            break
        successors = set(graph.adj[node])
        successors.difference_update(visited)
        pending.append(successors)
        visited.update(successors)

        data_out_queue.put_nowait(set_led(node, 0, 0x00F))
        for successor in successors:
            data_out_queue.put_nowait(set_led(successor, 0, 0xFF0))
        guess = await step()

        data_out_queue.put_nowait(set_led(node, 0, 0xFFF))
    await step()
    for node in visited:
        data_out_queue.put_nowait(set_led(node, 0, 0))


async def dfs_algorithm(start_node, data_out_queue):
    graph = get_graph()
    path = [(start_node, [])]

    #data_out_queue.put_nowait(set_led(target_node, 0, 0x800))
    while path:
        node, visited_successors = path[-1]
        algorithm_logger.info("visit %d", node)
        if node == target_node:
            data_out_queue.put_nowait(set_led(node, 0, 0x0F0))
            break
        if node in graph.adj:
            successors = set(graph.adj[node])
        else:
            successors = set()
        successors.difference_update(node for node, _ in path)
        successors.difference_update(visited_successors)

        data_out_queue.put_nowait(set_led(node, 0, 0x00F))
        for successor in successors:
            data_out_queue.put_nowait(set_led(successor, 0, 0xFF0))
        guess = await step()

        if not successors:
            del path[-1]
            data_out_queue.put_nowait(set_led(node, 0, 0))
        else:
            if guess in successors:
                path.append((guess, []))
            else:
                path.append((random.choice(tuple(successors)), []))
            visited_successors.append(path[-1][0])
            for successor in successors:
                if successor != path[-1][0]:
                    data_out_queue.put_nowait(set_led(successor, 0, 0))

            data_out_queue.put_nowait(set_led(node, 0, 0xFFF))
    await step()
    for node, _ in path:
        data_out_queue.put_nowait(set_led(node, 0, 0))



run_algorithm = dfs_algorithm
autostep = False
step_queue: Optional[asyncio.Queue] = None
running_algorithm: Optional[asyncio.Task] = None
button_times = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))

async def step():
    if autostep:
        await asyncio.sleep(5)
        return None
    else:
        return await step_queue.get()


async def blink_led(data_out_queue, remote_sender, led, colour):
    while True:
        data_out_queue.put_nowait(set_led(remote_sender, led, colour))
        await asyncio.sleep(3)
        data_out_queue.put_nowait(set_led(remote_sender, led, 0))
        await asyncio.sleep(3)


async def button_pressed(data_in_queue, data_out_queue, link_sender, remote_sender, buttons):
    global running_algorithm
    global step_queue
    global pairing_node_pin

    current_time = time.monotonic_ns()
    for button, state in enumerate(buttons):
        if not state:
            continue
        if button_times[remote_sender][button] + BUTTON_DEBOUNCE > current_time:
            processing_logger.info("button ignored: %d too close to %d", button_times[remote_sender][button], current_time)
            continue
        button_times[remote_sender][button] = current_time
        global step_queue
        if button == 0:
            if running_algorithm is not None and not running_algorithm.done():
                algorithm_logger.info("step from button")
                step_queue.put_nowait(remote_sender)
            else:
                algorithm_logger.info("start from button")
                step_queue = asyncio.Queue()
                running_algorithm = asyncio.create_task(run_algorithm(remote_sender, data_out_queue))
        else:
            if not connect_with_buttons:
                return
            pin = button - 1
            if pairing_node_pin is None:
                possible_colours = set(EDGE_COLOURS)
                remove_edge = None
                for u, v, d in buttons_graph.edges.data():
                    possible_colours.remove(d["colour"])
                    if d.get(remote_sender, {}).get("pin") == pin:
                        remove_edge = u, v, d
                        break
                if remove_edge is not None:
                    buttons_graph.remove_edge(remove_edge[0], remove_edge[1])
                    processing_logger.info("edge removed between %d:%d and %d:%d", remove_edge[0], remove_edge[2][remove_edge[0]]["pin"], remove_edge[1], remove_edge[2][remove_edge[1]]["pin"])
                    data_out_queue.put_nowait(set_led(remove_edge[0], remove_edge[2][remove_edge[0]]["pin"] + 1, 0))
                    data_out_queue.put_nowait(set_led(remove_edge[1], remove_edge[2][remove_edge[1]]["pin"] + 1, 0))
                    continue
                colour = random.choice(list(possible_colours))
                pairing_node_pin = remote_sender, pin, colour, asyncio.create_task(blink_led(data_out_queue, remote_sender, button, colour))
                processing_logger.info("started pairing %x:%d", remote_sender, pin)
            else:
                processing_logger.info("end pairing %x:%d", remote_sender, pin)
                pairing_node_pin[3].cancel()
                try:
                    await pairing_node_pin[3]
                except asyncio.CancelledError:
                    pass

                viable = True
                if remote_sender == pairing_node_pin[0]:
                    processing_logger.debug("pairing cancelled due to self loop")
                    viable = False
                if buttons_graph.has_edge(remote_sender, pairing_node_pin[0]):
                    processing_logger.warning("pairing cancelled due to duplicate edge")
                    viable = False
                if remote_sender in buttons_graph:
                    for _, _, d in buttons_graph.edges(remote_sender, data=True):
                        if d[remote_sender]["pin"] == pin:
                            processing_logger.warning("pairing cancelled due to pin in use")
                            viable = False

                if viable:
                    buttons_graph.add_edge(remote_sender, pairing_node_pin[0], colour=pairing_node_pin[2])
                    buttons_graph.edges[remote_sender, pairing_node_pin[0]][pairing_node_pin[0]] = {"pin": pairing_node_pin[1]}
                    buttons_graph.edges[remote_sender, pairing_node_pin[0]][remote_sender] = {"pin": pin}
                    data_out_queue.put_nowait(set_led(remote_sender, button, pairing_node_pin[2]))
                    data_out_queue.put_nowait(set_led(pairing_node_pin[0], pairing_node_pin[1] + 1, pairing_node_pin[2]))
                else:
                    data_out_queue.put_nowait(set_led(pairing_node_pin[0], pairing_node_pin[1] + 1, 0))
                processing_logger.info("new graph: %r", buttons_graph)
                pairing_node_pin = None


async def handle_message(n, data_in_queue, data_out_queue):
    message_type = n & 0b1100
    if message_type == 0b0000:
        # update neighbours
        link_sender = n >> 28
        remote_sender = (n >> 24) & 0xF
        neighbours = ((n >> 4) & 0xF, (n >> 8) & 0xF, (n >> 12) & 0xF, (n >> 16) & 0xF, (n >> 20) & 0xF)
        await update_neighbours(link_sender, remote_sender, neighbours)
    elif message_type == 0b0100:
        link_sender = n >> 28
        remote_sender = (n >> 24) & 0xF
        buttons = ((n >> 16) & 0b1, (n >> 17) & 0b1, (n >> 18) & 0b1, (n >> 19) & 0b1, (n >> 20) & 0b1, (n >> 21) & 0b1,
                   (n >> 22) & 0b1, (n >> 23) & 0b1)
        mapped_buttons = tuple(buttons[BUTTON_MAP[remote_sender, i]] for i in range(6))
        await button_pressed(data_in_queue, data_out_queue, link_sender, remote_sender, mapped_buttons)


async def message_handler(data_in_queue, data_out_queue):
    while True:
        n = await data_in_queue.get()
        await handle_message(n, data_in_queue, data_out_queue)


def reader(f, raw_in_queue):
    raw_in_queue.put_nowait(os.read(f, 1))

async def writer(f, raw_out_queue, writer_available):
    while True:
        await writer_available.wait()
        data = await raw_out_queue.get()

        writer_available.clear()
        os.write(f, data)

        raw_out_queue.task_done()

def notify_writer_available(writer_available):
    writer_available.set()


def _approach(target_value, old, step):
    if abs(target_value - old) < step:
        return target_value
    if target_value > old:
        return old + step
    return old - step

def _renderer(q):
    plt.ion()
    fig = plt.figure()
    pos = {}
    graph = None
    min_x, max_x, min_y, max_y, = -1, 1, -1, 1

    while True:
        fig.canvas.flush_events()
        try:
            graph = q.get_nowait()
        except queue.Empty:
            pass
        if graph is not None:
            fig.clear()
            pos = {n: pos[n] for n in graph.nodes() if n in pos}
            pos = networkx.arf_layout(graph, pos=pos, max_iter=1)
            step_ratio = 100
            if pos:
                target_min_x = min(x for (x, y) in pos.values()) - .1
                target_max_x = max(x for (x, y) in pos.values()) + .1
                target_min_y = min(y for (x, y) in pos.values()) - .1
                target_max_y = max(y for (x, y) in pos.values()) + .1
                x_step = (target_max_x - target_min_x) / step_ratio
                y_step = (target_max_y - target_min_y) / step_ratio
                min_x = _approach(target_min_x, min_x, x_step)
                max_x = _approach(target_max_x, max_x, x_step)
                min_y = _approach(target_min_y, min_y, y_step)
                max_y = _approach(target_max_y, max_y, y_step)
            ax = fig.add_subplot()
            ax.set_xlim(min_x, max_x)
            ax.set_ylim(min_y, max_y)
            labels = {n: str(n) for n in graph.nodes()}
            networkx.draw(graph, pos=pos, ax=ax, labels=labels)
            plt.draw()
            plt.show(block=False)
            fig.canvas.draw_idle()

async def renderer():
    q = multiprocessing.Queue(1)
    p = multiprocessing.Process(target=_renderer, args=(q,))
    p.start()
    try:
        while True:
            await asyncio.to_thread(q.put, get_graph())
            await asyncio.sleep(0.01)
    finally:
        p.terminate()


async def main(file):
    raw_in_queue = asyncio.Queue()
    raw_out_queue = asyncio.Queue()
    data_in_queue = asyncio.Queue()
    data_out_queue = asyncio.Queue()

    tasks = []
    if file == "":
        f = None
    else:
        f = os.open(sys.argv[1], os.O_RDWR)
    try:
        if f is not None:
            writer_available = asyncio.Event()
            asyncio.get_running_loop().add_reader(f, reader, f, raw_in_queue)
            asyncio.get_running_loop().add_writer(f, notify_writer_available, writer_available)
            tasks.append(writer(f, raw_out_queue, writer_available))
        else:
            tasks.append(pump_fake_data(raw_in_queue, raw_out_queue))

        tasks += [
            transceive_loop(raw_in_queue, raw_out_queue, data_in_queue, data_out_queue),
            message_handler(data_in_queue, data_out_queue),
            renderer(),
            update_neighbour_leds_loop(data_out_queue),
            refresh_neighbour_leds_loop(data_out_queue),
            send_neighbours_mode(data_out_queue),
        ]
        await asyncio.wait((asyncio.create_task(task) for task in tasks))
    finally:
        if f is not None:
            asyncio.get_running_loop().remove_reader(f)
            asyncio.get_running_loop().remove_writer(f)
            os.close(f)

async def pump_fake_data(raw_in_queue, raw_out_queue):
    d = b''
    raw_in_queue.put_nowait(b"\x00\x00\x00\x00")
    while True:
        m = await raw_out_queue.get()
        await asyncio.sleep(0.01)
        d += m
        try:
            di = d.index(b"\xFF\xFF\xFF\xFF\x00\x00\x00\x00")
        except ValueError:
            pass
        else:
            if len(d) >= di + 12:
                m = d[di+8:di+12]
                d = d[di+12:]
                n, = struct.unpack("<I", m)
                if n.bit_count() % 2:
                    if n & 0b10:
                        if random.random() < 0.1:
                            raw_in_queue.put_nowait(b"\x05\x00\x01\x11")
                        raw_in_queue.put_nowait(b"\x00\x02\x00\x11")
                        raw_in_queue.put_nowait(b"\x03\x01\x00\x12")
                    else:
                        raw_in_queue.put_nowait(b"\x00\x00\x00\x00")
        raw_out_queue.task_done()


if __name__ == "__main__":
    asyncio.run(main(sys.argv[1]))
